import torch
import matplotlib.pyplot as plt


def tanh(x):
    a = torch.exp(x) - torch.exp(-x)
    b = torch.exp(x) + torch.exp(-x)
    return a / b


x = torch.linspace(0, 1, 20)
y = torch.zeros(x.shape)
y[10:] = 1.

plt.plot(x, y, 'ro')

w1 = torch.tensor(0.1, requires_grad=True)
b1 = torch.tensor(0.2, requires_grad=True)
w2 = torch.tensor(0.1, requires_grad=True)
b2 = torch.tensor(0.2, requires_grad=True)
y_predict = tanh(w2 * (tanh(x * w1) + b1)) + b2
line, = plt.plot(x.detach().numpy(), y_predict.detach().numpy(), 'g^--')

epochs = 1000
for epoch in range(epochs):
    y_predict = tanh(w2 * (tanh(x * w1) + b1)) + b2
    line.set_data(x.detach().numpy(), y_predict.detach().numpy())
    e = torch.mean((y_predict - y) ** 2)
    e.backward()
    with torch.no_grad():
        w1 -= w1.grad * 0.05
        b1 -= b1.grad * 0.05
        w2 -= w1.grad * 0.05
        b2 -= b1.grad * 0.05
        w1.grad.zero_()
        b1.grad.zero_()
        w2.grad.zero_()
        b2.grad.zero_()
    print(e)
    plt.pause(0.1)
